use std::time::{Duration};
use tonic::{Code, Request, Status};
use tonic::transport::{Channel, Endpoint};
use tracing::{error, info, warn};

pub mod echo {
    tonic::include_proto!("echo");
}

use echo::{
    EchoRequest, EchoResponse, HealthCheckRequest, HealthCheckResponse,
    echo_service_client::EchoServiceClient,
};

#[derive(Debug)]
pub struct RetryPolicy {
    pub max_attempts: u32,
    pub initial_delay: Duration,
    pub max_delay: Duration,
    pub backoff_multiplier: f64,
    pub retryable_status_codes: Vec<Code>,
}

impl RetryPolicy {
    pub fn new() -> Self {
        Self {
            max_attempts: 5,
            initial_delay: Duration::from_secs(1),
            max_delay: Duration::from_secs(30),
            backoff_multiplier: 1.5,
            retryable_status_codes: vec![
                Code::Unavailable,
                Code::Unknown,
                Code::DeadlineExceeded,
            ],
        }
    }

    pub fn calculate_delay(&self, attempt: u32) -> Duration {
        let delay_secs =
            self.initial_delay.as_secs_f64() * self.backoff_multiplier.powi(attempt as i32 - 1);
        let delay_secs = delay_secs.min(self.max_delay.as_secs_f64());
        Duration::from_secs_f64(delay_secs)
    }
}

#[derive(Debug)]
pub struct GrpcConnectionManager {
    server_addr: String,
    retry_policy: RetryPolicy,
    broken_pipe_count: u32,
    successful_reconnects: u32,
    cancel_token: tokio_util::sync::CancellationToken,
}

impl GrpcConnectionManager {
    pub fn new(server_addr: String) -> Self {
        Self {
            server_addr,
            retry_policy: RetryPolicy::new(),
            broken_pipe_count: 0,
            successful_reconnects: 0,
            cancel_token: tokio_util::sync::CancellationToken::new(),
        }
    }

    pub async fn connect_with_retry(
        &mut self,
    ) -> Result<EchoServiceClient<tonic::transport::Channel>, Status> {
        let mut attempt = 0;

        while attempt < self.retry_policy.max_attempts {
            match EchoServiceClient::connect(self.server_addr.clone()).await {
                Ok(client) => {
                    info!("gRPC连接建立成功");
                    return Ok(client);
                }
                Err(e) => {
                    attempt += 1;

                    if e.to_string().contains("broken pipe") {
                        self.broken_pipe_count += 1;
                        warn!("检测到Broken pipe，准备重连...");
                    }

                    if attempt < self.retry_policy.max_attempts {
                        let delay = self.retry_policy.calculate_delay(attempt);
                        info!("第 {} 次重连尝试，等待 {:?}", attempt, delay);

                        tokio::time::sleep(delay).await;
                    }
                }
            }
        }

        Err(Status::unavailable("无法连接到gRPC服务器"))
    }

    pub async fn send_echo_with_recovery(
        &mut self,
        client: &mut EchoServiceClient<tonic::transport::Channel>,
        message: &str,
        sequence: i32,
        max_retries: u32,
    ) -> Result<EchoResponse, Status> {
        let request = EchoRequest {
            message: message.to_string(),
            sequence,
        };

        let mut retry_count = 0;

        loop {
            // 检查取消状态
            if self.cancel_token.is_cancelled() {
                warn!("操作取消");
                return Err(Status::cancelled("操作已取消"));
            }

            match client.echo(Request::new(request.clone())).await {
                Ok(response) => {
                    info!("Echo调用成功");
                    return Ok(response.into_inner());
                }
                Err(e) => {
                    // 正确的错误检测方式
                    // let is_broken_pipe = match e.code() {
                    //     Code::Unavailable => true,
                    //     Code::Unknown => e.message().contains("broken pipe"),
                    //     _ => false,
                    // };
                    let is_broken_pipe = self.should_reconnect(&e);

                    if is_broken_pipe && retry_count < max_retries {
                        retry_count += 1;
                        warn!("检测到Broken pipe错误，第{}次重连...", retry_count);

                        let new_client = self.connect_with_retry().await?;
                        *client = new_client;
                        self.successful_reconnects += 1;

                        info!("重连成功，重新发送消息...");
                        continue;
                    } else {
                        return Err(e);
                    }
                }
            }
        }
    }

    pub async fn test_connection_health(
        &self,
        client: &mut EchoServiceClient<tonic::transport::Channel>,
    ) -> Result<bool, Status> {
        let request = HealthCheckRequest {
            service: "echo".to_string(),
        };

        match client.health_check(Request::new(request)).await {
            Ok(_) => Ok(true),
            Err(_) => Ok(false),
        }
    }

    pub fn get_stats(&self) -> (u32, u32) {
        (self.broken_pipe_count, self.successful_reconnects)
    }

    pub fn get_cancel_token(&self) -> tokio_util::sync::CancellationToken {
        self.cancel_token.child_token()
    }

    pub async fn force_connection_reset(&mut self) -> Result<(), Box<dyn std::error::Error>> {
        info!("强制重置连接状态...");
        self.cancel_token.cancel();
        // 创建新的取消令牌用于后续操作
        self.cancel_token = tokio_util::sync::CancellationToken::new();
        Ok(())
    }

    pub async fn simulate_network_failure(&self) -> Result<(), Box<dyn std::error::Error>> {
        info!("模拟网络故障场景...");
        // 这里可以添加更复杂的网络故障模拟逻辑
        Ok(())
    }

    pub fn should_reconnect(&self, error: &(dyn std::error::Error + 'static)) -> bool {
        let error_str = error.to_string();
        error_str.contains("broken pipe") ||
            error_str.contains("connection reset") ||
            error_str.contains("EPIPE") ||
            error_str.contains("unavailable") ||
            error_str.contains("cancelled")
    }

    pub async fn simulate_broken_pipe(
        &mut self,
        client: &mut EchoServiceClient<Channel>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        info!("🎭 模拟Broken pipe错误场景");

        // 方法1: 创建无效的Endpoint触发连接错误
        let invalid_endpoint = Endpoint::from_shared("http://invalid-address:9999".to_string())?;

        // 短暂使用无效连接
        let _invalid_client = EchoServiceClient::connect(invalid_endpoint).await?;

        // 方法2: 主动关闭底层连接
        // 通过创建一个新的客户端来"重置"连接状态
        *client = self.connect_with_retry().await?;

        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt::init();

    let server_addr = "http://[::1]:50051".to_string();
    let mut connection_manager = GrpcConnectionManager::new(server_addr);

    info!("gRPC客户端启动，尝试连接到服务器...");

    let mut client = connection_manager.connect_with_retry().await?;

    // 模拟多次消息发送
    for i in 0..15 {
        let message = format!("Hello gRPC Server! Message #{}", i + 1);

        match connection_manager
            .send_echo_with_recovery(&mut client, &message, i as i32, 10)
            .await
        {
            Ok(response) => {
                info!(
                    "收到响应: {}, 服务器: {}",
                    response.message, response.server_id
                );

                // 随机模拟Broken pipe错误
                if i > 0 && i % 5 == 0 {
                    info!("模拟Broken pipe错误...");
                    // 短暂延迟让错误有机会传播
                    // tokio::time::sleep(Duration::from_millis(3000)).await;
                    //
                    // // 这里会触发重连逻辑
                    // // 方法1: 使用取消令牌主动取消当前操作
                    // let cancel_token = connection_manager.get_cancel_token();
                    // cancel_token.cancel();

                    if let Err(e) = connection_manager.simulate_broken_pipe(&mut client).await {
                        error!("❌ 模拟Broken pipe失败: {}", e);
                    }
                    // // 方法2: 强制重置连接状态
                    // connection_manager.force_connection_reset().await?;
                    //
                    // // 方法3: 模拟网络中断
                    // connection_manager.simulate_network_failure().await?;

                }
            }
            Err(e) => {
                error!("处理消息 #{} 时出错: {}", i + 1, e);
            }
        }

        tokio::time::sleep(Duration::from_secs(1)).await;
    }

    let (broken_pipes, successful_reconnects) = connection_manager.get_stats();
    info!(
        "客户端运行完成 - Broken pipes: {}, 成功重连: {}",
        broken_pipes, successful_reconnects
    );

    Ok(())
}
